package cn.fz.deepspeekdemo01.service;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Service
@RequiredArgsConstructor
public class DeepSeekService {
    @Autowired
    private OllamaChatModel ollamaChatModel;

    /**
     * 使用 ConcurrentHashMap 维护会话上下文（线程安全）
     */
    private final Map<String, List<Message>> chatContexts = new ConcurrentHashMap<>();

    public Flux<String> handleChatWebSocket(String sessionId,String question) {
        return sendChatRequest(sessionId, question)
                .doOnNext(response -> log.info("handleChatWebSocket 发送响应：{}", response))
                .onErrorResume(e->{
                    log.error("handleChatWebSocket WebSocket通信错误: {}", e.getMessage());
                    return Flux.just("抱歉, 服务器处理请求时发生错误, 请稍后重试.");
                });
    }

    public Flux<String> sendChatRequest(String sessionId,String question) {
        // 获取或创建会话历史
        List<Message> context = chatContexts
                .computeIfAbsent(sessionId, k -> new ArrayList<>());

        // 添加用户新消息到上下文
        context.add(new UserMessage(question));

        // 创建包含历史上下文的 Prompt
        ChatResponse response = ollamaChatModel.call(new Prompt(context));

        // 获取模型回复
        String answer = response.getResult().getOutput().getText();

        // 添加助手回复到上下文
        context.add(new AssistantMessage(answer));

        log.info("会话 {} 的上下文长度: {}", sessionId, context.size());

        return Flux.just(answer);
    }
}
